/**
 * 
 */
package org.swing.utility.common.stream;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;

import org.swing.utility.common.bytes.ByteBufferFifo;

/**
 * @author lqnhu
 *
 */
public class NIOReadStream {
	public NIOReadStream(ReadableByteChannel channel) {
		this.channel = channel;
	}

	/** @returns the number of bytes available to be read. */
	public int dataAvailable() {
		return totalAvailable;
	}

	public int getInt() {
		// TODO: Optimize?
		byte[] intbytes = new byte[4];
		getBytes(intbytes);
		int output = (((int) intbytes[0]) & 0xff);
		output |= (((int) intbytes[1]) & 0xff) << 8;
		output |= (((int) intbytes[2]) & 0xff) << 16;
		output |= (((int) intbytes[3]) & 0xff) << 24;
		return output;
	}

	public int getIntBigEndian() {
		// TODO: Optimize?
		byte[] intbytes = new byte[4];
		getBytes(intbytes);
		int output = (((int) intbytes[0]) & 0xff) << 24;
		output |= (((int) intbytes[1]) & 0xff) << 16;
		output |= (((int) intbytes[2]) & 0xff) << 8;
		output |= (((int) intbytes[3]) & 0xff);
		return output;
	}

	public void getBytes(byte[] output) {
		getBytes(output, 0, output.length);
	}

	public void getBytes(byte[] output, int offset, int length) {
		if (totalAvailable < length) {
			throw new IllegalStateException("Requested " + length
					+ " bytes; only have " + totalAvailable
					+ " bytes; call tryRead() first");
		}
		int originalOffset = offset;
		while (length > 0) {
			ByteBuffer first = buffers.getReadBuffer();
			assert first.remaining() > 0;
			// Copy bytes from first into output
			int bytesToCopy = length;
			if (first.remaining() < bytesToCopy)
				bytesToCopy = first.remaining();
			first.get(output, offset, bytesToCopy);
			offset += bytesToCopy;
			length -= bytesToCopy;
			assert offset <= output.length;
			assert length >= 0;
		}
		int bytesCopied = offset - originalOffset;
		totalAvailable -= bytesCopied;
	}

	public void close() {
		try {
			channel.close();
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
		buffers.clear();
	}

	/**
	 * Reads all available data from input. Returns true if data was read or the
	 * connection is still open. False indicates no data was read and the
	 * connection is closed.
	 *
	 * This allows the application an opportunity to consume the last bytes in
	 * the stream, without getting stuck if the connection is closed with a
	 * partial message.
	 */
	// TODO: Remove tryRead()?
	public boolean readAllAvailable() {
		// Read until it doesn't fill the buffer: indicates no more data
		// available
		boolean first = true;
		int lastRead = -1;
		while (true) {
			ByteBuffer writeBuffer = buffers.getWriteBuffer();
			assert writeBuffer.remaining() > 0;
			try {
				lastRead = channel.read(writeBuffer);
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
			if (lastRead > 0) {
				totalAvailable += lastRead;
			}
			if (writeBuffer.remaining() > 0) {
				break;
			}
			assert lastRead > 0;
			first = false;
		}
		if (lastRead == -1 && first) {
			// first read returned -1: connection is closed and buffer probably
			// has incomplete data
			return false;
		}
		return true;
	}

	/**
	 * Reads until we have at least desiredAvailable bytes buffered, there is no
	 * more data, or the channel is closed.
	 * 
	 * @returns number of bytes available for reading, or -1 if less than
	 *          desiredAvailable bytes are available, and the channel is closed.
	 */
	public int tryRead(int desiredAvailable) {
		// Read until we have enough, or read returns 0 or -1
		int lastRead = 1;
		while (lastRead > 0 && totalAvailable < desiredAvailable) {
			ByteBuffer writeBuffer = buffers.getWriteBuffer();
			assert writeBuffer.remaining() > 0;
			try {
				lastRead = channel.read(writeBuffer);
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
			if (lastRead > 0) {
				totalAvailable += lastRead;
			}
		}
		if (totalAvailable < desiredAvailable && lastRead == -1) {
			return -1;
		}
		return totalAvailable;
	}

	private final ReadableByteChannel channel;
	private final ByteBufferFifo buffers = new ByteBufferFifo();
	private int totalAvailable = 0;
}
